'''
Created on Sep 14, 2011

@author: oabalbin
'''
import numpy as np
import array
from optparse import OptionParser
from collections import defaultdict,deque #, OrderedDict
from bx.intervals.intersection import Interval, IntervalTree


class snpObject():
    def __init__(self,nsamples):
        self.nfields=4
        self.sample_array=np.zeros(nsamples*self.nfields)
        self.name=set()
        self.snptype=""
        self.DPREF_TOTAL=0
        self.DPALT_TOTAL=0
        self.NPATIENTS=np.zeros(nsamples)
    '''
    def __call__(self,nsamples):
        nfields=4
        self.sample_array=np.zeros(nsamples*nfields)
        self.name=set()
    '''

    def input_data(self,sample,data,gene_name,snptype):
        DPREF,DPALT,VFRAQ,HET=0,1,2,3
        HET_TH=0.80
        thsp=sample
        sample=self.nfields*sample
        self.sample_array[sample+DPREF]=data[DPREF]
        self.sample_array[sample+DPALT]=data[DPALT]
        self.sample_array[sample+VFRAQ]=float(data[DPALT])/(data[DPREF]+data[DPALT]) # variant fraction
        self.sample_array[sample+HET] = 0 if self.sample_array[sample+VFRAQ] < HET_TH else 1
        self.name.add(gene_name)
        self.snptype=snptype
        self.DPREF_TOTAL+=data[DPREF]
        self.DPALT_TOTAL+=+data[DPALT]
        self.NPATIENTS[thsp]=1

    
    def print_mydata(self):
        HET_TH=0.80
        DPFRAC_TOTAL=float(self.DPALT_TOTAL)/(self.DPREF_TOTAL+self.DPALT_TOTAL)
        HET = 0 if DPFRAC_TOTAL < HET_TH else 1
        line = ",".join(map(str,list(self.sample_array))).replace(',','\t')
        return self.snptype,",".join(self.name),line,str(self.DPREF_TOTAL),str(self.DPALT_TOTAL),str(DPFRAC_TOTAL),str(HET),str(np.sum(self.NPATIENTS))


class Gene:
    def __init__(self,start, end, value):
        self.start=start
        self.end=end
        self.value=value


def actionable_genes(gene_list):
    '''
    Creates a interval tree with the actionable genes
    '''
    
    actionable = IntervalTree()
    for gen in gene_list:
        start, end, obj = gen.start, gen.end, gen.value
        actionable.insert_interval( Interval(start, end, value=obj) )
    
    return actionable
        
        
def read_gene_list(ifile):
    '''
    Read a bed file with fields: chr, start, end, name, mark, strand
    '''
    ifile=open(ifile)
    actionable_genes=deque()
    for l in ifile:
        fields = l.strip('\n').split('\t')
        start, end, value=int(fields[1]),int(fields[2]), {'chr':fields[0], 'annot':fields[3],'strand':fields[5]}
        actionable_genes.append( Gene(start, end, value) )
    
    return actionable_genes
    
def mark_actionable_genes(ifile,snp_list):
    '''
    Input a list of actionable genes, and a list of total snps. 
    '''
    gene_list = read_gene_list(ifile)
    act_genes = actionable_genes(gene_list)
    act_snps = defaultdict()
    
    for snp in snp_list:
        loc = snp.split('@')
        #print snp, loc
        chr,pos=loc[0],int(loc[1].split('|')[0])
        chr,pos=loc[0],int(loc[1])
        genes = act_genes.find(pos,pos)
        
        for g in genes:
            if chr == g.value['chr']:
                act_snps[snp] = g.value['annot']
    
    return act_snps


def snp_type(info_field):
    '''
    Determine the type of snp.
    '''        
    gene_annot={'NON_SYNONYMOUS_CODING':1,'STOP_GAINED':2,'STOP_LOST':3,
                 'NON_SYNONYMOUS_CODING(SPLICE_SITE)':4,'STOP_GAINED(SPLICE_SITE)':5,
                 'ESSENTIAL_SPLICE_SITE':6, 'SPLICE_SITE':7,
                 'SYNONYMOUS_CODING':8,'3\'UTR':9,'5\'UTR':10,'INTRONIC':11,'INTERGENIC':12,
                 'DOWSTREAM':13,'UPSTREAM':14}
    
    # The part below handles the snp annotation and the gene name associated to it. 
    
    info_items=defaultdict()
    for s in info_field:
        k,i=s.split('=')
        info_items[k]=i

    snv_annotation = set(gene_annot.keys()).intersection(set(info_items.keys()))
    
    if snv_annotation:
        tmpid=100                   
        for gt in snv_annotation:
            id = gene_annot[gt]
            if id < tmpid:
                tmp,tmpid=gt,id
        snptype=tmpid #id
        gene_name=info_items[tmp]
    else:
        snptype=999
        gene_name='nan'
        tmp='nan'
    return [snptype,gene_name,tmp]


def write_snps(file,snpsOfInterest,nsamples):
    '''
    write the snps to an output file
    '''
    ofile =open(file,'w')
    bfile = open(file+'.bed','w')
    pfile = open(file+'polyp','w')
    sample_tags = []
    tags = ["DPREF","DPALT","VFRAQ","HOM"]

    for sp in nsamples:
        sample_tags+= [t+'_'+sp for t in tags]
    
    #h=["#CHROM","POS","REF","ALT","TYPE","ID"]+["DPREF","DPALT","VFRAQ","HET"]*nsamples+["DPREF_TOTAL","DPALT_TOTAL","VFRAQ_TOTAL","HOM"]
    h=["#CHROM","POS","REF","ALT","TYPE","ID"]+sample_tags+["DPREF_TOTAL","DPALT_TOTAL","VFRAQ_TOTAL","HOM","NPATIENTS"]
    
    ofile.write(",".join(h).replace(',','\t')+'\n')
    for snpid,snp  in snpsOfInterest.iteritems():
        
        ofile.write(snpid.replace('@','\t')+'\t'+",".join(list(snp.print_mydata())).replace(',','\t')+'\n')
        bed=[snpid.split('@')[0],str(int(snpid.split('@')[1])-1),str(int(snpid.split('@')[1])+1)]
        
        polyp=[snpid.split('@')[0]+':'+snpid.split('@')[1],snpid.split('@')[2]+'/'+snpid.split('@')[3]]
        bfile.write(",".join(bed).replace(',','\t')+'\n')
        pfile.write(",".join(polyp).replace(',','\t')+'\n')
        
    ofile.close()
    bfile.close()

def read_vcf_file(files_list, ofile,genesOfinterest):
    '''
    Read a vcf file format 4.0 produce by SamTools    
    '''
    
    nfiles=len(files_list)
    snpsOfInterest = defaultdict()
    non_syn_snps = 7
    sample_names = ['c']*nfiles
    
    for n, ifile in enumerate(files_list):
        print n, ifile.split('/')[-1]
        sample_names[n]=str(n)+'_'+ifile.split('.')[1]
        
        ifile = open(ifile)
        
        for l in ifile:
            if l.startswith('#'):
                continue
            fields=l.strip('\n').replace(',','$').split('\t')
            chr,pos,qual,ref,alt=fields[0],fields[1],fields[5],fields[3],fields[4]
            info = fields[7].split(';')
            # Do not consider reference homozygous or INDELs
            if alt==".":
                continue
            if info[0]=="INDEL":
                # If the indel is inside the genes of interest include it
                continue
            
            # Filter Out SNPs syno
            snptype,db_gene_name,snptype_name=snp_type(info)
            if snptype > non_syn_snps:
                continue
            
            # Check if the snp position is inside genes of interest
            snpid=chr+'@'+str(pos)+'@'+ref+'@'+alt
             
            genes = genesOfinterest.find(int(pos),int(pos))
            
            for g in genes:
                if chr == g.value['chr']: # The snps is in the chromosome of interest. So it is in the gene of interesT
                    if snpid not in snpsOfInterest.keys():
                        snpOD = snpObject(nfiles)
                        snpsOfInterest[snpid]=snpOD
  
                    dpc = info[3].split('=')[1]
                    ref_support, alt_support = int(dpc.split('$')[0]) + int(dpc.split('$')[1]),int(dpc.split('$')[2]) + int(dpc.split('$')[3])
                    map_qual=info[4].split("=")[1]
                    # Inputing the info into the snpObject
                    data=np.array([ref_support,alt_support])
                    gene_name=g.value['annot']+'@'+db_gene_name

                    snpsOfInterest[snpid].input_data(n,data,gene_name,snptype_name)
                    #Sprint snpid,snptype_name,snpsOfInterest[snpid].print_mydata()

        ifile.close()
    
    write_snps(ofile,snpsOfInterest,sample_names)

if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] ")
    # Annotation files
    optionparser.add_option("-a", "--genesOfInterest", dest="genesOfInterest",
                            help="genesOfInterest genes bed file")
    optionparser.add_option("-f", "--vcf_files",action="append", type="str", dest="vcf_files",
                        help="vcf_files. ")
    optionparser.add_option("-o", "--ofile", dest="ofile",
                        help="outputfile")

    
    (options, args) = optionparser.parse_args()    


    gfile=options.genesOfInterest #'/exds/users/oabalbin/projects/exomes/nunez2/chron_relevant_genelist_hg19.bed'
    files_list=options.vcf_files #['/exds/users/oabalbin/projects/exomes/nunez2/sam_calls_quick/nunez3788benign.sam_snps.mpileup_annot_isec_exomePad50h.vcf','/exds/users/oabalbin/projects/exomes/nunez2/sam_calls_quick/nunez3787benign.sam_snps.mpileup_annot_isec_exomePad50h.vcf']

    ofile = options.ofile # '/exds/users/oabalbin/projects/exomes/nunez2/sam_calls_quick/test.tmp'
    gene_list = read_gene_list(gfile)
    genesOfinterest = actionable_genes(gene_list)
    read_vcf_file(files_list,ofile,genesOfinterest)
